import torch
import torch.nn as nn

from Modules.SoftMax2dArea import SoftMax2dArea

class MapConv(nn.Module):
    def __init__(self, in_channel, hidden_channel=64, adapt_pool_size=7):
        super().__init__()
        self.in_channel = in_channel
        self.hidden_channel = hidden_channel
        self.adapt_pool_size = adapt_pool_size
        self.map_conv = nn.Sequential(
            nn.Conv2d(in_channel, hidden_channel, kernel_size=1, stride=1, padding=0),
            nn.LeakyReLU(inplace=True),
            nn.AdaptiveAvgPool2d(adapt_pool_size),
            nn.Conv2d(hidden_channel, 1, kernel_size=1),
            SoftMax2dArea()
        )


    def forward(self, x):
        return self.map_conv(x)


if __name__ == '__main__':
    pass


